#if !defined(MBEDTLS_CONFIG_FILE)
#include "mbedtls/config.h"
#else
#include MBEDTLS_CONFIG_FILE
#endif

#if defined(MBEDTLS_SM2_C)

#include "ssm2.h"
#include "mbedtls/entropy.h"
#include "mbedtls/ctr_drbg.h"

#if defined(MBEDTLS_PLATFORM_C)
#include "mbedtls/platform.h"
#else
#include <stdio.h>
#include <stdlib.h>

#define mbedtls_calloc      calloc
#define mbedtls_printf      printf
#define mbedtls_free        free
#endif


static int sm2_get_rand(const mbedtls_ecp_group *grp, mbedtls_mpi *r, int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
{
	int ret;
	size_t blind_tries = 0;
	do
	{
		size_t nlen = (grp->nbits + 7) / 8;
		MBEDTLS_MPI_CHK(mbedtls_mpi_fill_random(r, nlen, f_rng, p_rng));
		MBEDTLS_MPI_CHK(mbedtls_mpi_shift_r(r, 8 * nlen - grp->nbits));

		/* See mbedtls_ecp_gen_keypair() */
		if (++blind_tries > 30)
			return (MBEDTLS_ERR_SM2_RANDOM_FAILED);
	}
	while (mbedtls_mpi_cmp_int(r, 1) < 0 ||
	mbedtls_mpi_cmp_mpi(r, &grp->N) >= 0);
cleanup:
	return (ret);
}


//秘钥生成
int mbedtls_ssm2_genkey_part1(mbedtls_ecp_group *grp,
	const mbedtls_ecp_point *G,
	mbedtls_mpi *d, mbedtls_ecp_point *Q,
	int(*f_rng)(void *, unsigned char *, size_t),
	void *p_rng)
{
	int ret;

	// d1∈[1，n-1]
	MBEDTLS_MPI_CHK(sm2_get_rand(grp, d, f_rng, p_rng));

	// p1＝d1[*]G
	// Multiplication R = m * P
	ret = mbedtls_ecp_mul(grp, Q, d, G, f_rng, p_rng);
cleanup:
	return ret;
}

int mbedtls_ssm2_genkey_part2(mbedtls_ecp_group *grp,
	const mbedtls_ecp_point *G,
	const mbedtls_ecp_point *Q1,
	mbedtls_mpi *d, mbedtls_ecp_point *Q,
	int(*f_rng)(void *, unsigned char *, size_t),
	void *p_rng)
{
	int ret;
	mbedtls_mpi d_inv, d_one;

	mbedtls_mpi_init(&d_inv);
	mbedtls_mpi_init(&d_one);

	// d1∈[1，n-1]
	MBEDTLS_MPI_CHK(sm2_get_rand(grp, d, f_rng, p_rng));

	// p＝d2^-1[*]p1[-]G = d_inv[*]p1[-]G
	// d_inv = d2^-1
	//* Modular inverse: X = A^-1 mod N  (HAC 14.61 / 14.64)
	ret = mbedtls_mpi_inv_mod(&d_inv, d, &grp->N);
	if (ret != 0)
		goto cleanup;

	// points by integers: R = m * P + n * Q
	// int mbedtls_ecp_muladd( mbedtls_ecp_group *grp, mbedtls_ecp_point *R,
	//           const mbedtls_mpi *m, const mbedtls_ecp_point *P,
	//           const mbedtls_mpi *n, const mbedtls_ecp_point *Q );
	// d_inv [*] G [-] G

	mbedtls_mpi_lset(&d_one, -1);
	ret = mbedtls_ecp_muladd(grp, Q, &d_inv, Q1, &d_one, G);

cleanup:
	mbedtls_mpi_free(&d_inv);
	mbedtls_mpi_free(&d_one);
	return ret;
}


//解密
int mbedtls_ssm2_decrypt_step1(mbedtls_ecp_group *grp, const mbedtls_mpi *d1,
    const mbedtls_ecp_point *C1, mbedtls_ecp_point *T1)
{
  int ret;
  //T1＝D1[*]C1
  ret = mbedtls_ecp_mul(grp, T1, d1, C1, NULL, NULL);
  return ret;
}
int mbedtls_ssm2_decrypt_step2(mbedtls_ecp_group *grp, const mbedtls_mpi *d2,
    const mbedtls_ecp_point *T1, mbedtls_ecp_point *T2)
{
  int ret;
  mbedtls_mpi inv;
  mbedtls_mpi_init(&inv);
  //T2＝D2^-1[*]T1
  ret = mbedtls_mpi_inv_mod(&inv, d2, &grp->N);
  if (ret!=0)
    goto cleanup;
  ret = mbedtls_ecp_mul(grp, T2, &inv, T1, NULL, NULL);
cleanup:
  mbedtls_mpi_free(&inv);
  return ret;
}
int mbedtls_ssm2_decrypt_step3(mbedtls_ecp_group *grp, const mbedtls_ecp_point *C1,
    const mbedtls_ecp_point *T2, mbedtls_ecp_point *P)
{
  int ret;

  mbedtls_mpi one;
  mbedtls_mpi done;
  mbedtls_mpi_init(&one);
  mbedtls_mpi_init(&done);

  // Q = (x2，y2)＝T2[-]C1
  //   = (D2^-1*T1)-C1
  //   = (D2^-1 * D1*C1) [-]C1
  //   = ((D2^-1)*D1-1) [-] C1
  // D = (D2^-1)*D1-1

  mbedtls_mpi_lset(&one, 1);
  mbedtls_mpi_lset(&done, -1);

  ret = mbedtls_ecp_muladd(grp, P, &one, T2, &done, C1 );
  if (ret!=0)
  {
    goto cleanup;
  }

cleanup:
  mbedtls_mpi_free(&one);
  mbedtls_mpi_free(&done);

  return ret;
}

//签名
int mbedtls_ssm2_sign_step1(const mbedtls_ecp_group *grp,
    mbedtls_mpi *k1, mbedtls_ecp_point *q1,
    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
{
  int ret;

  // k1∈[1，n-1]
  MBEDTLS_MPI_CHK( sm2_get_rand(grp,  k1, f_rng, p_rng ) );

  // R1＝k1[*]G
  ret = mbedtls_ecp_mul( (mbedtls_ecp_group*)grp, q1, k1, &grp->G, f_rng, p_rng );
cleanup:
  return ret;
}
int mbedtls_ssm2_sign_step2(
    mbedtls_ecp_group *grp,
    const mbedtls_mpi *d2,        
    const mbedtls_ecp_point *q1, 
    const mbedtls_mpi *e,        
    mbedtls_mpi *r,
    mbedtls_mpi *s2,
    mbedtls_mpi *s3,
    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
{
  int	ret = -1;

  mbedtls_mpi one;
  mbedtls_mpi k2;
  mbedtls_mpi k3;

  mbedtls_ecp_point q2;

  mbedtls_mpi_init(&k2);
  mbedtls_mpi_init(&k3);
  mbedtls_mpi_init(&one);
  mbedtls_mpi_lset(&one, 1);
  mbedtls_ecp_point_init(&q2);

  // k2∈[1，n-1]
  MBEDTLS_MPI_CHK( sm2_get_rand(grp, &k2, f_rng, p_rng ) );

  // q2＝k2[*]G
  ret = mbedtls_ecp_mul( grp, &q2, &k2, &grp->G, f_rng, p_rng );
  if(ret!=0)
    goto cleanup;

  do {
    // k3∈[1，n-1]
    MBEDTLS_MPI_CHK( sm2_get_rand(grp, &k3, f_rng, p_rng ) );
    // (x1，y1)＝k3[*]q1[+]q2
    // points by integers: R = m * P + n * Q
    // int mbedtls_ecp_muladd( mbedtls_ecp_group *grp, mbedtls_ecp_point *R,
    //           const mbedtls_mpi *m, const mbedtls_ecp_point *P,
    //           const mbedtls_mpi *n, const mbedtls_ecp_point *Q );

    ret = mbedtls_ecp_muladd( grp, &q2, &k3, q1, &one, &q2 );
    if(ret!=0)
      goto cleanup;
    // r＝x1+e mod n
    ret = mbedtls_mpi_add_mpi( r, &q2.X, e );
    if (ret!=0)
      goto cleanup;
    ret = mbedtls_mpi_mod_mpi( r, r, &grp->N );
    if (ret!=0)
      goto cleanup;
  } while (mbedtls_mpi_cmp_int(r, 0) == 0); //r != 0

  // s2＝d2*k3 mod n
  ret = mbedtls_mpi_mul_mpi(s2, d2, &k3);
  if (ret!=0)
    goto cleanup;
  ret = mbedtls_mpi_mod_mpi(s2, s2, &grp->N);
  if (ret!=0)
    goto cleanup;

  // s3＝d2*(r+k2) mod n
  ret = mbedtls_mpi_add_mpi(s3, r, &k2);
  if (ret!=0)
    goto cleanup;
  ret = mbedtls_mpi_mod_mpi(s3, s3, &grp->N);
  if (ret!=0)
    goto cleanup;
  ret = mbedtls_mpi_mul_mpi(s3, s3, d2);
  if (ret!=0)
    goto cleanup;
  ret = mbedtls_mpi_mod_mpi(s3, s3, &grp->N);
cleanup:

  mbedtls_mpi_free( &one );
  mbedtls_mpi_free( &k2 );
  mbedtls_mpi_free( &k3 );

  mbedtls_ecp_point_free( &q2 );

  return ret;
}
int mbedtls_ssm2_sign_step3(const mbedtls_ecp_group* grp,
    const mbedtls_mpi *d1, const mbedtls_mpi *k1,
    const mbedtls_mpi *r, const mbedtls_mpi *s2, const mbedtls_mpi *s3,
    mbedtls_mpi *s)
{
  int ret;
  mbedtls_mpi d_inv;
  mbedtls_mpi t_i1;
  mbedtls_mpi t_i2;

  mbedtls_mpi_init(&d_inv);
  mbedtls_mpi_init(&t_i1);
  mbedtls_mpi_init(&t_i2);

  // s＝(d1^-1*k1)s2+d1^-1*s3-r mod n
  //  = d_inv * k1 * s2 + d_inv * s3 -r mod n

  // d_inv = d1^-1
  //* Modular inverse: X = A^-1 mod N  (HAC 14.61 / 14.64)
  ret = mbedtls_mpi_inv_mod(&d_inv, d1, &grp->N);
  if (ret!=0)
    goto cleanup;

  // t_i1 = d_inv * k1 * s2
  ret = mbedtls_mpi_mul_mpi(&t_i1, &d_inv, k1);
  if (ret!=0) {
    goto cleanup;
  }

  ret = mbedtls_mpi_mul_mpi(&t_i1, &t_i1, s2);
  if (ret!=0) {
    goto cleanup;
  }

  // t_i2 = d_inv * s3
  ret = mbedtls_mpi_mul_mpi(&t_i2, &d_inv, s3);
  if (ret!=0) {
    goto cleanup;
  }

  // t_i2 = t_i2 + t_i1
  ret = mbedtls_mpi_add_mpi(&t_i2, &t_i2, &t_i1);
  if (ret!=0) {
    goto cleanup;
  }

  // t_i2 = t_i2 - r
  ret = mbedtls_mpi_sub_mpi(&t_i2, &t_i2, r);
  if (ret!=0) {
    goto cleanup;
  }

  //s = t_i2 mod n
  ret = mbedtls_mpi_mod_mpi(s, &t_i2, &grp->N);
  if (ret!=0) {
    goto cleanup;
  }

  //s = 0 or s + r = n
  ret = mbedtls_mpi_add_mpi(&t_i2, s, r);
  if (ret!=0) {
    goto cleanup;
  }

  if( mbedtls_mpi_cmp_int( s, 0 ) == 0 || mbedtls_mpi_cmp_mpi(&t_i2, &grp->N)==0)
  {
    ret = -2;
    goto cleanup;
  }

  ret = 0;
cleanup:

  mbedtls_mpi_free(&d_inv);
  mbedtls_mpi_free(&t_i1);
  mbedtls_mpi_free(&t_i2);

  return ret;
}


//秘钥协商
int mbedtls_ssm2_computekey_step1(
    mbedtls_ecp_group *grp,
    const mbedtls_mpi *d1,
    mbedtls_mpi *k1,
    mbedtls_mpi *R1,
    int (*f_rng)(void *, unsigned char *, size_t),
    void *p_rng)
{
  int ret;

  // d1∈[1，n-1]
  MBEDTLS_MPI_CHK( sm2_get_rand(grp, k1, f_rng, p_rng ) );

  MBEDTLS_MPI_CHK( mbedtls_mpi_mul_mpi(R1, k1, d1 ) );

cleanup:
  return ret;
}

int mbedtls_ssm2_computekey_step2(
    mbedtls_ecp_group *grp,
    const mbedtls_mpi *d2,
    const mbedtls_mpi *R1,
    mbedtls_mpi *k2,
    mbedtls_ecp_point *R,
    int (*f_rng)(void *, unsigned char *, size_t),
    void *p_rng)
{
  int ret;

  mbedtls_mpi t;
  mbedtls_mpi_init(&t);

  // d1∈[1，n-1]
  MBEDTLS_MPI_CHK( sm2_get_rand(grp, k2, f_rng, p_rng ) );
  MBEDTLS_MPI_CHK( mbedtls_mpi_inv_mod(&t, d2, &grp->N) );    //d2^-1
  MBEDTLS_MPI_CHK( mbedtls_mpi_mul_mpi(&t, &t, R1 ) );
  MBEDTLS_MPI_CHK( mbedtls_mpi_mul_mpi(&t, &t, k2 ) );
  MBEDTLS_MPI_CHK( mbedtls_mpi_mod_mpi(&t, &t, &grp->N ) );
  MBEDTLS_MPI_CHK( mbedtls_ecp_mul(grp, R, &t, &grp->G, NULL, NULL));

cleanup:
  mbedtls_mpi_free(&t);
  return ret;
}

int mbedtls_ssm2_computekey_step3(
    mbedtls_ecp_group *grp,
    const mbedtls_ecp_point *Ra,    //本端临时公钥 Ra or Rb  x1bar
//    const mbedtls_ecp_point *Pa,    //本端加密公钥 Pa or Pb
    const mbedtls_ecp_point *Rb,    //对端临时公钥 Rb or Ra
//    const mbedtls_ecp_point *Pb,    //对端加密公钥 Pb or Pa
    mbedtls_mpi* x1bar,
    mbedtls_mpi* x2bar)
{
  int ret;
  int w;
  mbedtls_mpi two_pow_w;
  mbedtls_mpi one;
  mbedtls_mpi_init(&two_pow_w);
  mbedtls_mpi_init(&one);

  mbedtls_mpi_lset(&one, 1);
  /*
   * w = ceil(keybits / 2) - 1
   * x = 2^w + (x and (2^w - 1)) = 2^w + (x mod 2^w)
   */
  w = (mbedtls_mpi_size(&grp->N) + 1) / 2 - 1;            //向上取整 w = 127
  mbedtls_mpi_lset(&two_pow_w, 1);
  ret = mbedtls_mpi_shift_l(&two_pow_w, w);               //two_pow_w = 2^w
  if (ret!=0)
  {
    goto cleanup;
  }

  //Ra x1bar
  //Rb x2bar
  //Caculate self x1bar
  ret = mbedtls_mpi_mod_mpi(x1bar, &Ra->X, &two_pow_w);
  if (ret!=0)
  {
    goto cleanup;
  }
  ret = mbedtls_mpi_add_mpi(x1bar, x1bar, &two_pow_w);
  if (ret!=0)
  {
    goto cleanup;
  }

  //Caculate peer x2bar
  ret = mbedtls_mpi_mod_mpi(x2bar, &Rb->X, &two_pow_w);
  if (ret!=0)
  {
    goto cleanup;
  }
  ret = mbedtls_mpi_add_mpi(x2bar, x2bar, &two_pow_w);
  if(ret!=0)
  {
    goto cleanup;
  }
cleanup:
  mbedtls_mpi_free(&two_pow_w);
  mbedtls_mpi_free(&one);
  return ret;
}

int mbedtls_ssm2_computekey_step4(
    const mbedtls_mpi* r1,
    const mbedtls_mpi* x1bar,
    mbedtls_mpi* i)
{
  return mbedtls_mpi_mul_mpi(i, r1, x1bar);     // xx1*r1
}

int mbedtls_ssm2_computekey_step5(
    mbedtls_ecp_group *grp,
    const mbedtls_mpi* r2,
    const mbedtls_mpi* dA2,
    const mbedtls_mpi* i_from_client,
    mbedtls_mpi* i)
{
  int ret;
  mbedtls_mpi dd2;
  mbedtls_mpi_init(&dd2);

  MBEDTLS_MPI_CHK(mbedtls_mpi_inv_mod(&dd2, dA2, &grp->N)); // d2^-1
  MBEDTLS_MPI_CHK(mbedtls_mpi_mul_mpi(i, i_from_client, r2));         // 2a. xx1*r1*r2
  MBEDTLS_MPI_CHK(mbedtls_mpi_add_int(i, i, 1));          // 2b
  MBEDTLS_MPI_CHK(mbedtls_mpi_mul_mpi(i, i, &dd2));       // 2c
cleanup:
  mbedtls_mpi_free(&dd2);
  return ret;
}

int mbedtls_ssm2_computekey_step6(
    mbedtls_ecp_group *grp,
    const mbedtls_mpi *dA1,
    const mbedtls_mpi *i_from_server,
    mbedtls_mpi *t)
{
  int ret;
  MBEDTLS_MPI_CHK(mbedtls_mpi_mul_mpi(t, i_from_server, dA1));
  MBEDTLS_MPI_CHK(mbedtls_mpi_sub_int(t, t, 1));

  MBEDTLS_MPI_CHK(mbedtls_mpi_mod_mpi(t, t, &grp->N));
cleanup:
  return ret;
}

#endif
